"""Module to implement boundary conditions."""

class BoundaryCondition(object):
    """Base class for all boundary conditions"""

    def apply_bc(self, grid):
        """Do nothing. This is the default behavior for all solvers."""
        pass

class PeriodicBC(BoundaryCondition):
    r"""Periodic boundary conditions in 1D as described in LeVeque (section
    7.1)

    For a grid with :math:`n_b` boundary cells (:math:`n_b = 2` for a
    five point stencil), the periodic boundary conditions is
    implemented by setting:

    .. math::

        Q_1 = Q_{-3} \,\, , Q_0 = Q_{-4}

        Q_{-2} = Q_2 \,\, , Q_{-1} = Q_3

    """
    def apply_bc(self, grid):

        cells = grid.boundary_cells["periodic"]["cells"]
        images = grid.boundary_cells["periodic"]["images"]

        for i, cell in enumerate(cells):
            image = images[i]

            cell.q[:] = image.q[:]

class SolidWallLeft1D(BoundaryCondition):
    r"""Solid wall or reflecting boundary conditions as described by LeVeque
    (section 7.3.3)

    Reflecting boundary conditions are implemented by symmetry. All
    physical variables are extrapolated normally while the velocity is
    reflected.

    For the implementation sake, if you have a system of conservation
    laws in which the :math:`j`th conserved variable deals with
    velocity, then you need to provide this index at the time of
    construction. This is so that I know which conserved variable to
    simply extrapolate and which to flip (velocity)

    """
    def __init__(self, velocity):
        """Constructor.

        Parameters:
        ------------

        velocity : int
            Index for the velocity variable in the vector of
            conserved variables.

        """

        self.velocity = velocity

    def apply_bc(self, grid):
        """Apply the boundary condition."""
        q = grid.q
        nvar = grid.nvar
        
        q1 = q[:, 1]; q2 = q[:,2]
        q0 = q[:, 0]; q3 = q[:,3]

        q1[:] = q2
        q0[:] = q3

        q1[self.velocity] *= -1
        q0[self.velocity] *= -1

class SolidWallRight1D(BoundaryCondition):
    r"""Solid wall or reflecting boundary conditions as described by LeVeque
    (section 7.3.3)

    Reflecting boundary conditions are implemented by symmetry. All
    physical variables are extrapolated normally while the velocity is
    reflected.

    For the implementation sake, if you have a system of conservation
    laws in which the :math:`j`th conserved variable deals with
    velocity, then you need to provide this index at the time of
    construction. This is so that I know which conserved variable to
    simply extrapolate and which to flip (velocity)

    """
    def __init__(self, velocity):
        """Constructor.

        Parameters:
        ------------

        velocity : int
            Index for the velocity variable in the vector of
            conserved variables.

        """

        self.velocity = velocity

    def apply_bc(self, grid):
        """Apply the boundary condition."""
        q = grid.q
        nvar = grid.nvar

        q0 = q[:, -2]; q3 = q[:,-3]
        q1 = q[:, -1]; q2 = q[:,-4]

        q0[:] = q3
        q1[:] = q2

        q0[self.velocity] *= -1
        q1[self.velocity] *= -1

class TransmissiveLeft1D(BoundaryCondition):
    r"""Simple boundary condition so that waves leave the solution domain
    cleanly.

    The boundary conditions are implemented by simply extending the
    physical domain values to the ghost cells.

    """
    def apply_bc(self, grid):
        """Apply the boundary condition."""
        q = grid.q; nb = grid.nb
        for i in range(nb):
            q[:, i] = q[:, nb]

class TransmissiveRight1D(BoundaryCondition):
    r"""Simple boundary condition so that waves leave the solution domain
    cleanly.

    The boundary conditions are implemented by simply extending the
    physical domain values to the ghost cells.

    """
    def apply_bc(self, grid):
        """Apply the boundary condition."""
        q = grid.q; nb = grid.nb
        for i in range(nb):
            q[:, -(i+1)] = q[:, -(nb+1)]
